from typing import Dict
from common.cache.pool_cache import KeyPoolCache
from common.utils.tools import ToolsUtil


class VideoOpenAiHkServer:
    """
    Luma

    Document:
        https://openai-hk.com/docs/lab/luma-video-v2.html
    """

    def __init__(self):
        self.baseUrl = "https://api.openai-hk.com"

    @classmethod
    def check_params(cls, params: Dict[str, any]):
        """
        验证请求参数

        Args:
            params (Dict[str, any]): 请求参数

        Author:
            zero
        """
        if params.get("scale") not in ["1:1", "3:4", "4:3", "9:16", "16:9"]:
            raise Exception("视频比例不在合规值内: [1:1, 3:4, ...]")
        if not params.get("prompt") and not params.get("image_url"):
            raise Exception("视频描述和图片至少填一项")

    async def submit(self, params: dict):
        """
        发起生成请求

        Args:
            params (Dict[str, any]): 请求参数

        Return:
            {
                'id': 'd335c8c7-bce0-4df0-9f5d-da1d658fc5cc',
                'prompt': '',
                'state': 'pending',
                'last_frame': None,
                'thumbnail': None,
                'video': None,
                'video_raw': None
            }

        Author:
            zero
        """
        # 验证参数
        self.check_params(params)

        # 获取密钥
        token = await self.get_apikey()

        # 提交任务
        try:
            other_params = params.get("other_params") or {}
            response = await ToolsUtil.post_curl(
                url=f"{self.baseUrl}/luma/generations",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer " + str(token)
                },
                data={
                    # 提示词
                    "user_prompt": params.get("prompt"),
                    # 视频宽高比
                    "aspect_ratio": params.get("scale"),
                    # 关键帧: 开始图片
                    "image_url": params.get("image_url", ""),
                    # 关键帧: 结束图片
                    "image_end_url": str(other_params.get("image_end_url", "")),
                    # 是否扩展提示词
                    "expand_prompt": bool(other_params.get("expand_prompt") or 0),
                    # 是否循环尾呼应
                    "loop": bool(other_params.get("loop") or 0),
                    # 模型名称: [ray-v1, ray-v2]
                    "model_name": params.get("model", "ray-v2")
                }
            )

            # 结果处理
            if response.get("error"):
                code = response["error"]["code"]
                msgs = response["error"]["message"]
                raise Exception(f"{code}: {msgs}")

            # 结果返回
            return response
        except Exception as e:
            raise Exception(str(e))

    async def fetch(self, task_id: str):
        """
        获取生成结果

        Returns:
            dict

        Author:
            zero
        """
        token: str = await self.get_apikey()
        try:
            response = await ToolsUtil.get_curl(
                url=f"{self.baseUrl}/luma/v2/generations/{task_id}",
                headers={
                    "Content-Type": str("application/json"),
                    "Authorization": "Bearer " + str(token)
                }
            )
            return response
        except Exception as e:
            raise Exception(str(e))

    @classmethod
    async def get_apikey(cls) -> str:
        """
        获取密钥

        Returns:
            str

        Author:
            zero
        """
        apikey = await KeyPoolCache(scene="video", channel="luma-hk").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise Exception("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]


class VideoZhipuServer:
    """
    智谱AI视频

    Document:
        https://bigmodel.cn/dev/api/videomodel/cogvideox
    """

    def __init__(self):
        self.baseUrl = "https://open.bigmodel.cn/api/paas/v4"

    @classmethod
    def check_params(cls, params: Dict[str, any]):
        """
        验证请求参数

        Args:
            params (Dict[str, any]): 请求参数

        Author:
            zero
        """
        other_params = params.get("other_params") or {}
        if params.get("model") not in ["cogvideox-2", "cogvideox-flash"]:
            raise Exception("视频生成模型异常")
        if not params.get("prompt") and not params.get("image_url"):
            raise Exception("视频描述和图片至少填一项")
        if int(other_params.get("fps", 30)) not in [30, 60]:
            raise Exception("视频帧率不在合规值内: [30, 60]")
        if other_params.get("quality") not in ["speed", "quality"]:
            raise Exception("视频模式不在合规值内: [speed, quality]")
        if params.get("scale") not in [
            "720x480", "1024x1024", "1280x960", "960x1280",
            "1920x1080", "1080x1920", "2048x1080", "3840x2160"
        ]:
            raise Exception("视频比例不在合规值内")

    async def submit(self, params: Dict[str, any]):
        """
        发起生成请求

        Args:
            params (Dict[str, any]): 请求参数

        Returns:
            {
                'id': '1251721982564066-8835116647168602944',
                'model': 'cogvideox-2',
                'request_id': '20250401222341dbd08e5d2d4e4127',
                'task_status': 'PROCESSING'
            }

        Author:
            zero
        """
        # 验证参数
        self.check_params(params)

        # 获取密钥
        token: str = await self.get_apikey()

        try:
            other_params = params.get("other_params") or {}
            response = await ToolsUtil.post_curl(
                url=f"{self.baseUrl}/videos/generations",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer " + str(token)
                },
                data={
                    # 模型编码: [cogvideox-2, cogvideox-flash]
                    "model": params.get("model", "cogvideox-2"),
                    # 视频描述
                    "prompt": params.get("prompt", ""),
                    # 视频尺寸
                    "size": params.get("scale", "1920x1080"),
                    # 参考图像
                    "image_url": params.get("image_url", ""),
                    # AI音效
                    "with_audio": bool(other_params.get("with_audio", False)),
                    # 输出模式: [speed=速度优先, quality=质量优先]
                    "quality": str(other_params.get("quality", "speed")),
                    # 视频帧率: [30, 60]
                    "fps": int(other_params.get("fps", 30))
                }
            )

            if response.get("error"):
                code = response["error"]["code"]
                msgs = response["error"]["message"]
                raise Exception(f"{code}: {msgs}")

            return response
        except Exception as e:
            raise Exception(str(e))

    async def fetch(self, task_id: str) -> dict:
        """
        获取生成结果

        Returns:
            # 成功
            {
                'model': 'cogvideox-2',
                'request_id': '20250401222341dbd08e5d2d4e4127',
                'task_status': 'SUCCESS',
                'video_result': [{
                    'cover_image_url': 'https://aigc-files.bigmodel.cn/api/cogvideo/e913c01a-0f0_cover_0.jpeg',
                    'url': 'https://aigc-files.bigmodel.cn/api/cogvideo/e913c01a-0f04-11f0-ad3e-d2a032528fb8_0.mp4'
                }]
            }
            # 失败
            {'error': {'code': '1000', 'message': '身份验证失败。'}}

        Author:
            zero
        """
        token: str = await self.get_apikey()
        try:
            response = await ToolsUtil.get_curl(
                url=f"{self.baseUrl}/async-result/{task_id}",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer " + str(token)
                }
            )
            return response
        except Exception as e:
            raise Exception(str(e))

    @classmethod
    async def get_apikey(cls) -> str:
        """
        获取密钥

        Returns:
            str

        Author:
            zero
        """
        apikey = await KeyPoolCache(scene="video", channel="zhipu").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise Exception("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]
